import os
import glob

# 设置路径A
path_A = r'C:\Users\Caibin\Desktop\mixup_20250123_seg\mixup\val\labels'  # 替换为你的txt文件夹路径


def process_file(file_path):
    with open(file_path, 'r') as f:
        lines = f.readlines()

    new_lines = []
    for line in lines:
        parts = line.strip().split()

        # 确保至少有 20 个元素（类别 + 4 个框参数 + 6 对关键点 + 6 个可见性标记）
        if len(parts) < 20:
            continue

        # 保留类别
        new_line = [parts[0]]

        # 跳过 center_x、center_y、width、height
        keypoints = parts[5:]  # 从第6个开始都是关键点信息

        # 每3个一组（x, y, visible），只保留 x 和 y
        for i in range(0, len(keypoints), 3):
            new_line.append(keypoints[i])  # x
            new_line.append(keypoints[i + 1])  # y

        # 组合并加入新行
        new_lines.append(" ".join(new_line))

    # 覆盖写入原文件
    with open(file_path, 'w') as f:
        f.write("\n".join(new_lines))


def main():
    # 查找路径A下的所有txt文件
    txt_files = glob.glob(os.path.join(path_A, "*.txt"))

    # 对每个txt文件进行处理
    for txt_file in txt_files:
        process_file(txt_file)
        print(f"已处理: {txt_file}")


if __name__ == "__main__":
    main()